import h5py
import numpy as np
import torch
from ResNet import resnet50,resnet152,resnet18
from mixer import MLPMixer
with h5py.File('/home/zsc/ADS_B_select_files_/dataset/Preprocess_64x64_19.h5','r') as f:
    X_train = np.array(f['X_train'][:])
    y_train = np.array(f['y_train'][:])
    X_test = np.array(f['X_test'][:])
    y_test = np.array(f['y_test'][:])
print(len(X_train),len(y_train))
print(len(X_test),len(y_test))

3405 3405
30645 30645

from DataSet import DataSet64x64,DataLoader,DataSet16x16
train_data = DataSet64x64(X_train, y_train)
test_data = DataSet64x64(X_test, y_test)

train_loader = DataLoader(train_data,batch_size=64,shuffle=True,num_workers=4)
test_loader = DataLoader(test_data,batch_size=64,shuffle=True,num_workers=4)

def update_lr(optimizer, lr):    
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
lr = 0.001
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
model = MLPMixer(
    image_size = 64,
    patch_size = 16,
    channels=1,
    dim = 512,
    depth = 2,
    num_classes = 34
).to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr = lr)
save_path = './model/mixer_19'
best_acc_train = 0
best_acc_test = 0
for epoch in range(50):
    model.train()
    running_loss = 0
    for i, (x,y) in enumerate(train_loader):
        X , Y = x.to(device),y.to(device)
        optimizer.zero_grad()
        outputs = model(X)
        loss = criterion(outputs,Y)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        if (i+1) % 100 == 0:
            print('Epoch {},  [ {} / {} ] , RunningLoss: {:.4f}'.format(epoch+1,i+1,len(train_loader),running_loss/100))
            running_loss = 0
    
    if (epoch+1) % 5 == 0:
            lr /= 2
            update_lr(optimizer, lr)
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for x,y in test_loader:
            x,y = x.to(device),y.to(device)
            outputs = model(x)
            _ , predicted = torch.max(outputs.data,1)
            total += y.size(0)
            correct += (predicted == y).sum().item()
        Tcorrect = 0
        Ttotal = 0
        for x,y in train_loader:
            x,y = x.to(device),y.to(device)
            outputs = model(x)
            _ , predicted = torch.max(outputs.data,1)
            Ttotal += y.size(0)
            Tcorrect += (predicted == y).sum().item()
        print('Epoch {}, acc (TestSet): {:.4f}%, acc (TrainSet) : {:.4f}%'.format(epoch+1,correct/total*100,Tcorrect/Ttotal*100))

    if best_acc_train + best_acc_test <= Tcorrect/Ttotal + correct/total:
        torch.save(model.state_dict(),save_path)
        best_acc_test = Tcorrect/Ttotal
        best_acc_train = correct/total

cuda:0
Epoch 1, acc (TestSet): 5.6714%, acc (TrainSet) : 6.4317%
Epoch 2, acc (TestSet): 6.2131%, acc (TrainSet) : 6.7254%
Epoch 3, acc (TestSet): 10.2301%, acc (TrainSet) : 10.8370%
Epoch 4, acc (TestSet): 20.1501%, acc (TrainSet) : 21.7327%
Epoch 5, acc (TestSet): 60.3067%, acc (TrainSet) : 65.5800%
Epoch 6, acc (TestSet): 83.9517%, acc (TrainSet) : 93.2452%
Epoch 7, acc (TestSet): 86.9212%, acc (TrainSet) : 97.2100%
Epoch 8, acc (TestSet): 89.5513%, acc (TrainSet) : 98.8253%
Epoch 9, acc (TestSet): 90.8272%, acc (TrainSet) : 99.5888%
Epoch 10, acc (TestSet): 91.6006%, acc (TrainSet) : 99.4714%
Epoch 11, acc (TestSet): 91.7344%, acc (TrainSet) : 99.5888%
Epoch 12, acc (TestSet): 91.9269%, acc (TrainSet) : 99.8238%
Epoch 13, acc (TestSet): 92.0509%, acc (TrainSet) : 99.9413%
Epoch 14, acc (TestSet): 91.8551%, acc (TrainSet) : 99.7651%
Epoch 15, acc (TestSet): 92.0020%, acc (TrainSet) : 99.8825%
Epoch 16, acc (TestSet): 92.0901%, acc (TrainSet) : 100.0000%
Epoch 17, acc (TestSet): 92.2043%, acc (TrainSet) : 100.0000%
Epoch 18, acc (TestSet): 92.1977%, acc (TrainSet) : 100.0000%
Epoch 19, acc (TestSet): 92.2565%, acc (TrainSet) : 100.0000%
Epoch 20, acc (TestSet): 92.2565%, acc (TrainSet) : 100.0000%
Epoch 21, acc (TestSet): 92.2630%, acc (TrainSet) : 100.0000%
Epoch 22, acc (TestSet): 92.2793%, acc (TrainSet) : 100.0000%
Epoch 23, acc (TestSet): 92.3054%, acc (TrainSet) : 100.0000%
Epoch 24, acc (TestSet): 92.3185%, acc (TrainSet) : 100.0000%
Epoch 25, acc (TestSet): 92.3315%, acc (TrainSet) : 100.0000%
Epoch 26, acc (TestSet): 92.3348%, acc (TrainSet) : 100.0000%
Epoch 27, acc (TestSet): 92.3250%, acc (TrainSet) : 100.0000%
Epoch 28, acc (TestSet): 92.3479%, acc (TrainSet) : 100.0000%
Epoch 29, acc (TestSet): 92.3511%, acc (TrainSet) : 100.0000%
Epoch 30, acc (TestSet): 92.3674%, acc (TrainSet) : 100.0000%
Epoch 31, acc (TestSet): 92.3740%, acc (TrainSet) : 100.0000%
Epoch 32, acc (TestSet): 92.3446%, acc (TrainSet) : 100.0000%
Epoch 33, acc (TestSet): 92.3576%, acc (TrainSet) : 100.0000%
Epoch 34, acc (TestSet): 92.3642%, acc (TrainSet) : 100.0000%
Epoch 35, acc (TestSet): 92.3511%, acc (TrainSet) : 100.0000%
Epoch 36, acc (TestSet): 92.3609%, acc (TrainSet) : 100.0000%
Epoch 37, acc (TestSet): 92.3479%, acc (TrainSet) : 100.0000%
Epoch 38, acc (TestSet): 92.3576%, acc (TrainSet) : 100.0000%
Epoch 39, acc (TestSet): 92.3707%, acc (TrainSet) : 100.0000%
Epoch 40, acc (TestSet): 92.3674%, acc (TrainSet) : 100.0000%
Epoch 41, acc (TestSet): 92.3609%, acc (TrainSet) : 100.0000%
Epoch 42, acc (TestSet): 92.3772%, acc (TrainSet) : 100.0000%
Epoch 43, acc (TestSet): 92.3870%, acc (TrainSet) : 100.0000%
Epoch 44, acc (TestSet): 92.3772%, acc (TrainSet) : 100.0000%
Epoch 45, acc (TestSet): 92.3837%, acc (TrainSet) : 100.0000%
Epoch 46, acc (TestSet): 92.3903%, acc (TrainSet) : 100.0000%
Epoch 47, acc (TestSet): 92.3903%, acc (TrainSet) : 100.0000%
Epoch 48, acc (TestSet): 92.3837%, acc (TrainSet) : 100.0000%
Epoch 49, acc (TestSet): 92.3870%, acc (TrainSet) : 100.0000%
Epoch 50, acc (TestSet): 92.3870%, acc (TrainSet) : 100.0000%

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# model = resnet18(1,24).to(device)
model = MLPMixer(
    image_size = 64,
    patch_size = 16,
    channels=1,
    dim = 512,
    depth = 2,
    num_classes = 34
).to(device)
model.load_state_dict(torch.load('./model/mixer_19'))
X_f = []
y_f = []
mat = np.zeros((34,34))
with torch.no_grad():
    # for x,y in train_loader:
    #     x = x.to(device)
    #     outputs = model(x)
    #     for i,X in enumerate(outputs):
    #         X_f.append(X.cpu().detach().numpy())
    #         y_f.append(int(y[i].detach().numpy()))
    for x,y in test_loader:
        x = x.to(device)
        outputs = model(x)
        for i,X in enumerate(outputs):
            
            X_f.append(X.cpu().detach().numpy())
            y_f.append(int(y[i].detach().numpy()))
            _ , predicted = torch.max(X.unsqueeze(0),1)
            preL = int(predicted.cpu().detach().numpy())
            trueL = int(y[i].detach().numpy())
            mat[trueL,preL] += 1

linesum = mat.sum(1)
linesum.shape
mat1 = mat/mat.sum(axis=1 ,keepdims = True)

import seaborn as sns
import matplotlib.pyplot as plt

plt.figure(figsize=(20,20))
sns.heatmap(data=mat1,annot=True,cmap= "YlGnBu")

<AxesSubplot:>

{'780df9': 0, '78050f': 1, '780d19': 2, '7804bc': 3, '780830': 4, '780800': 5, '7802b1': 6, '780fe3': 7, '780b69': 8, '780070': 9, '780534': 10, '780d8c': 11, '78087a': 12, '780649': 13, '780fb6': 14, '7805c3': 15, '7804f5': 16, '7809a3': 17, '78027b': 18, '7804ea': 19, '780609': 20, '78027c': 21, '780277': 22, '780c2d': 23}

from Tsne import DecAndPlot

DecAndPlot(X_f,y_f) 

No handles with labels found to put in legend.